## 成交量均线 VMA（默认20周期），用于判断成交量是否显著放大：
## 条件：当前分钟成交量 > VMA * 倍数（默认1.5）。
## 数据传入方式参考 RSI：按tick增量更新，在内部按分钟聚合。
## 新增：滑动窗口VMA，确保只基于最近时间窗口计算，避免被早期高成交量污染

from collections import deque
from typing import Dict, Optional, Union, Iterable
import pandas as pd
from datetime import datetime, timedelta


class solution:

    def __init__(self, period: int = 20, factor: float = 1.5, min_periods: int = 5, window_minutes: int = 30):
        self.period: int = max(1, int(period))
        self.factor: float = float(factor)
        self.min_periods: int = max(1, int(min_periods))
        # 新增：滑动窗口时间限制（分钟），确保VMA只基于最近N分钟内的数据
        self.window_minutes: int = max(1, int(window_minutes))

        # 每只股票的分钟聚合与历史
        self._minute_key: Dict[str, Optional[str]] = {}
        self._minute_volume: Dict[str, int] = {}
        # 修改：存储带时间戳的成交量数据，用于滑动窗口过滤
        self._history_volumes: Dict[str, deque] = {}  # 存储 (timestamp, volume) 元组
        # 新增：存储分钟时间戳，用于滑动窗口过滤
        self._minute_timestamps: Dict[str, deque] = {}

    def reset(self, stock_code: Optional[str] = None) -> None:
        if stock_code is None:
            self._minute_key.clear()
            self._minute_volume.clear()
            self._history_volumes.clear()
            self._minute_timestamps.clear()
            return
        self._minute_key.pop(stock_code, None)
        self._minute_volume.pop(stock_code, None)
        self._history_volumes.pop(stock_code, None)
        self._minute_timestamps.pop(stock_code, None)

    def _to_minute_key(self, timestamp: Union[int, float, str, pd.Timestamp]) -> Optional[str]:
        try:
            if timestamp is None:
                return None
            if isinstance(timestamp, pd.Timestamp):
                dt = timestamp
            elif isinstance(timestamp, (int, float)):
                # 视为毫秒时间戳，转为北京时间
                dt = pd.to_datetime(timestamp, unit='ms') + pd.Timedelta(hours=8)
            else:
                s = str(timestamp)
                if s.isdigit():
                    # 形如 20250101103015 按上海时间解析
                    dt = pd.to_datetime(s, format='%Y%m%d%H%M%S', errors='coerce')
                else:
                    dt = pd.to_datetime(s, errors='coerce')
            if pd.isna(dt):
                return None
            return dt.strftime('%Y-%m-%d %H:%M')
        except Exception:
            return None

    def _cleanup_expired_data(self, stock_code: str) -> None:
        """清理过期的成交量数据，只保留滑动窗口内的数据"""
        try:
            if stock_code not in self._minute_timestamps or stock_code not in self._history_volumes:
                return

            timestamps = self._minute_timestamps[stock_code]
            volumes = self._history_volumes[stock_code]

            if not timestamps or not volumes:
                return

            # 计算滑动窗口的截止时间
            cutoff_time = datetime.now() - timedelta(minutes=self.window_minutes)

            # 从队列头部开始清理过期数据
            removed_count = 0
            while timestamps and volumes:
                if timestamps[0] < cutoff_time:
                    timestamps.popleft()
                    volumes.popleft()
                    removed_count += 1
                else:
                    break

            # 调试信息：如果清理了数据，打印日志
            if removed_count > 0:
                print(f"🧹 {stock_code} 清理过期数据: {removed_count} 个数据点")

            # 添加清理后的状态信息
            remaining_count = len(volumes) if volumes else 0
            print(
                f"📊 {stock_code} 清理后状态: 剩余{remaining_count}个数据点, 窗口截止时间={cutoff_time.strftime('%H:%M:%S')}")

        except Exception as e:
            print(f"❌ {stock_code} 清理过期数据异常: {e}")
            pass

    def _check_minute_roll(self, stock_code: str) -> None:
        """主动检查是否需要分钟切换，确保分钟数据及时更新"""
        try:
            if stock_code not in self._minute_key:
                return

            current_minute = datetime.now().strftime('%Y-%m-%d %H:%M')
            stored_minute = self._minute_key.get(stock_code)

            if stored_minute != current_minute:
                # 分钟已切换，需要写入历史
                print(f"🔄 {stock_code} 分钟切换: {stored_minute} -> {current_minute}")
                self._roll_minute(stock_code)
                self._minute_key[stock_code] = current_minute

        except Exception as e:
            print(f"❌ {stock_code} 检查分钟切换异常: {e}")

    def _roll_minute(self, stock_code: str) -> None:
        """在分钟切换时，把上一分钟的成交量写入历史队列。"""
        vol = int(self._minute_volume.get(stock_code, 0))
        if vol < 0:
            vol = 0

        if stock_code not in self._history_volumes:
            self._history_volumes[stock_code] = deque(maxlen=self.period * 2)  # 增大容量以容纳滑动窗口
        if stock_code not in self._minute_timestamps:
            self._minute_timestamps[stock_code] = deque(maxlen=self.period * 2)

        # 记录当前分钟的时间戳和成交量
        current_time = datetime.now()
        self._minute_timestamps[stock_code].append(current_time)
        self._history_volumes[stock_code].append(vol)

        # 不再清理数据，让数据自然增长
        # self._cleanup_expired_data(stock_code)

        # 重置当前分钟累计量
        self._minute_volume[stock_code] = 0

    def warmup_with_volumes(self, stock_code: str, minute_volumes: Iterable[Union[int, float]]) -> None:
        """
        直接以分钟成交量序列进行预热（仅写入已完成的分钟）。
        适用于在中途重启后，用当日已有的分钟线成交量来填充VMA历史。
        """
        if stock_code not in self._history_volumes:
            self._history_volumes[stock_code] = deque(maxlen=self.period * 2)
        if stock_code not in self._minute_timestamps:
            self._minute_timestamps[stock_code] = deque(maxlen=self.period * 2)

        dq_vol = self._history_volumes[stock_code]
        dq_time = self._minute_timestamps[stock_code]

        try:
            # 只取最近 window_minutes 内的数据，避免预热过多历史数据
            current_time = datetime.now()
            cutoff_time = current_time - timedelta(minutes=self.window_minutes)

            # 限制预热数据量，最多只预热 window_minutes 个数据点
            max_warmup_count = min(self.window_minutes, self.period)
            volumes_list = list(minute_volumes)[-max_warmup_count:]

            for i, v in enumerate(volumes_list):
                try:
                    iv = int(float(v))
                    if iv >= 0:
                        # 为每个预热数据分配时间戳，从当前时间往前推
                        timestamp = current_time - timedelta(minutes=len(volumes_list) - i)
                        dq_time.append(timestamp)
                        dq_vol.append(iv)
                except Exception:
                    continue

            # 不再清理数据，让数据自然增长
            # self._cleanup_expired_data(stock_code)

        except Exception:
            pass

    def warmup_from_minutes(self, stock_code: str, minutes_df: pd.DataFrame,
                            include_current_minute: bool = False) -> None:
        """
        使用分钟线DataFrame预热：需包含时间列与成交量列。
        - 自动识别时间列（包含'time'字样）与成交量列（'volume'/'vol'/'成交量'）。
        - 默认仅写入“已完成”的分钟；若 include_current_minute=True，则包含最后一条分钟。
        """
        try:
            if minutes_df is None or minutes_df.empty:
                return
            cols = [c for c in minutes_df.columns]
            time_col = next((c for c in cols if 'time' in c.lower()), None)
            if time_col is None:
                return
            vol_col_candidates = ['volume', 'vol', '成交量']
            vol_col = next((c for c in vol_col_candidates if c in minutes_df.columns), None)
            if vol_col is None:
                return
            df = minutes_df.copy()
            # 统一为datetime并按分钟排序
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')
            df = df[df['__time__'].notna()].sort_values('__time__')
            if df.empty:
                return
            if not include_current_minute:
                # 仅纳入严格早于“当前分钟开始”的数据
                now_floor = pd.to_datetime(datetime.now().strftime('%Y-%m-%d %H:%M'))
                df = df[df['__time__'] < now_floor]
                if df.empty:
                    return
            self.warmup_with_volumes(stock_code, df[vol_col].tolist())
        except Exception:
            pass

    def on_tick(self, stock_code: str, timestamp: Union[int, float, str, pd.Timestamp],
                volume: Optional[Union[int, float]]) -> None:
        """
        按tick增量更新：内部按分钟聚合成交量。
        参数中的 volume 建议传入该tick的成交量（手/张），若为None则忽略。
        """
        try:
            if volume is None:
                return
            vol = int(float(volume))
        except Exception:
            return

        minute_key = self._to_minute_key(timestamp)
        if minute_key is None:
            return

        prev_key = self._minute_key.get(stock_code)
        if prev_key is None:
            self._minute_key[stock_code] = minute_key
            self._minute_volume[stock_code] = vol
            if stock_code not in self._history_volumes:
                self._history_volumes[stock_code] = deque(maxlen=self.period * 2)
                self._minute_timestamps[stock_code] = deque(maxlen=self.period * 2)
            print(f"🔍 {stock_code} 初始化: 分钟={minute_key}, 成交量={vol}")
            return

        if prev_key == minute_key:
            old_vol = self._minute_volume.get(stock_code, 0)
            new_vol = old_vol + vol
            self._minute_volume[stock_code] = new_vol
            print(f"🔍 {stock_code} tick累计: 分钟={minute_key}, +{vol} -> {new_vol}")
            return

        # 分钟切换：写入历史并开始新的分钟累计
        print(f"🔄 {stock_code} 分钟切换: {prev_key} -> {minute_key}")
        self._roll_minute(stock_code)
        self._minute_key[stock_code] = minute_key
        self._minute_volume[stock_code] = vol
        print(f"🔍 {stock_code} 新分钟开始: 分钟={minute_key}, 成交量={vol}")

    def get_vma(self, stock_code: str) -> Optional[float]:
        """返回最近 period 个已完成分钟的平均成交量（VMA）"""
        # 不主动检查分钟切换，避免重置当前分钟累计量
        # self._check_minute_roll(stock_code)

        history = self._history_volumes.get(stock_code)
        if not history or len(history) < self.min_periods:
            return None

        try:
            # 只取最近 period 个数据计算VMA
            recent_volumes = list(history)[-self.period:]
            return float(sum(recent_volumes)) / float(len(recent_volumes))
        except Exception:
            return None

    def get_sliding_window_vma(self, stock_code: str) -> Optional[float]:
        """返回滑动窗口内所有数据的平均成交量（更精确的VMA）"""
        # 不主动检查分钟切换，避免重置当前分钟累计量
        # self._check_minute_roll(stock_code)

        try:
            history = self._history_volumes.get(stock_code)
            if not history:
                print(f"⚠️ {stock_code} 滑动窗口VMA: 无历史数据")
                return None

            if len(history) < self.min_periods:
                print(f"⚠️ {stock_code} 滑动窗口VMA: 数据点不足 {len(history)} < {self.min_periods}")
                return None

            # 计算滑动窗口的截止时间
            cutoff_time = datetime.now() - timedelta(minutes=self.window_minutes)

            # 动态筛选最近时间窗口内的数据，而不是清理数据
            timestamps = self._minute_timestamps.get(stock_code, [])
            if not timestamps:
                print(f"⚠️ {stock_code} 滑动窗口VMA: 无时间戳数据")
                return None

            # 找到在时间窗口内的数据
            window_volumes = []
            for i, (ts, vol) in enumerate(zip(timestamps, history)):
                if ts >= cutoff_time:
                    window_volumes.append(vol)
                # 调试：打印每个数据点的时间和信息
                print(f"  📅 {stock_code} 数据点{i}: 时间={ts.strftime('%H:%M:%S')}, 成交量={vol}")

            print(
                f"🔍 {stock_code} 滑动窗口VMA: 总数据点={len(history)}, 时间戳={len(timestamps)}, 窗口内数据点={len(window_volumes)}, 截止时间={cutoff_time.strftime('%H:%M:%S')}")

            if len(window_volumes) < self.min_periods:
                print(f"⚠️ {stock_code} 滑动窗口VMA: 窗口内数据点不足 {len(window_volumes)} < {self.min_periods}")
                return None

            result = float(sum(window_volumes)) / float(len(window_volumes))
            print(f"✅ {stock_code} 滑动窗口VMA计算成功: {result:.2f} (基于{len(window_volumes)}个数据点)")
            return result

        except Exception as e:
            print(f"❌ {stock_code} 计算滑动窗口VMA异常: {e}")
            return None

    def make_big(self, stock_code: str) -> bool:
        """
        判断当前分钟累计成交量是否显著放大：当前分钟量 > VMA * factor。
        若历史不足 min_periods 或无当前分钟量，则返回 False。
        """
        # 不主动检查分钟切换，避免重置当前分钟累计量
        # self._check_minute_roll(stock_code)

        current_vol = int(self._minute_volume.get(stock_code, 0))
        vma = self.get_vma(stock_code)
        if vma is None or current_vol <= 0:
            return False
        try:
            return float(current_vol) > float(vma) * float(self.factor)
        except Exception:
            return False

    def was_prev_minute_big(self, stock_code: str) -> bool:
        """
        判断上一已完成分钟是否放量：上一分钟量 > 之前分钟VMA * factor。
        要求：至少 min_periods+1 个已完成分钟（上一分钟 + 至少 min_periods 个历史）。
        """
        # 不主动检查分钟切换，避免重置当前分钟累计量
        # self._check_minute_roll(stock_code)

        dq = self._history_volumes.get(stock_code)
        if dq is None:
            return False
        try:
            history = list(dq)
            if len(history) < max(2, self.min_periods + 1):
                return False
            prev_vol = float(history[-1])
            base = history[:-1]
            base_tail = base[-self.period:]
            if not base_tail:
                return False
            base_avg = float(sum(base_tail)) / float(len(base_tail))
            if base_avg <= 0:
                return False
            return prev_vol > base_avg * float(self.factor)
        except Exception:
            return False

    def get_state(self, stock_code: str) -> Optional[dict]:
        """返回当前成交量状态，便于调试或展示。"""
        try:
            # 不主动检查分钟切换，避免重置当前分钟累计量
            # self._check_minute_roll(stock_code)

            try:
                current_vol = int(self._minute_volume.get(stock_code, 0))
            except Exception:
                current_vol = 0

            vma = self.get_vma(stock_code)
            sliding_vma = self.get_sliding_window_vma(stock_code)
            history = self._history_volumes.get(stock_code)
            history_len = len(history) if history is not None else 0

            # 获取时间窗口信息
            window_info = "无数据"
            if history and self._minute_timestamps.get(stock_code):
                timestamps = self._minute_timestamps[stock_code]
                if timestamps:
                    oldest_time = timestamps[0]
                    newest_time = timestamps[-1]
                    time_span = (newest_time - oldest_time).total_seconds() / 60
                    window_info = f"{oldest_time.strftime('%H:%M')}-{newest_time.strftime('%H:%M')} ({time_span:.1f}分钟)"

            return {
                'current_volume': current_vol,
                'vma': float(vma) if vma is not None else None,
                'sliding_vma': float(sliding_vma) if sliding_vma is not None else None,
                'factor': float(self.factor),
                'is_big': bool(self.make_big(stock_code)) if vma is not None else False,
                'history_len': history_len,
                'period': self.period,
                'window_minutes': self.window_minutes,
                'window_data_count': history_len,
                'window_time_span': window_info
            }
        except Exception as e:
            return None

    def set_params(self, period: Optional[int] = None, factor: Optional[float] = None,
                   min_periods: Optional[int] = None, window_minutes: Optional[int] = None) -> None:
        """运行时调整参数；如修改 period，会同步调整已有股票的历史队列长度。"""
        changed = False
        if period is not None and int(period) >= 1 and int(period) != self.period:
            self.period = int(period)
            changed = True
        if factor is not None:
            self.factor = float(factor)
        if min_periods is not None and int(min_periods) >= 1:
            self.min_periods = int(min_periods)
        if window_minutes is not None and int(window_minutes) >= 1:
            self.window_minutes = int(window_minutes)
            # 窗口时间改变时，不再清理数据，让数据自然增长
            # for code in list(self._history_volumes.keys()):
            #     self._cleanup_expired_data(code)

        if changed:
            for code, dq in list(self._history_volumes.items()):
                try:
                    items = list(dq)[-self.period:]
                    self._history_volumes[code] = deque(items, maxlen=self.period * 2)
                except Exception:
                    self._history_volumes[code] = deque(maxlen=self.period * 2)

    def debug_sliding_window(self, stock_code: str) -> None:
        """调试滑动窗口状态，打印详细信息"""
        try:
            print(f"\n🔍 {stock_code} 滑动窗口调试信息:")
            print(f"  窗口大小: {self.window_minutes} 分钟")
            print(f"  当前时间: {datetime.now().strftime('%H:%M:%S')}")

            if stock_code in self._minute_timestamps and stock_code in self._history_volumes:
                timestamps = self._minute_timestamps[stock_code]
                volumes = self._history_volumes[stock_code]

                if timestamps and volumes:
                    print(f"  历史数据点: {len(volumes)}")
                    print(f"  最早时间: {timestamps[0].strftime('%H:%M:%S')}")
                    print(f"  最新时间: {timestamps[-1].strftime('%H:%M:%S')}")

                    # 计算滑动窗口截止时间
                    cutoff_time = datetime.now() - timedelta(minutes=self.window_minutes)
                    print(f"  滑动窗口截止: {cutoff_time.strftime('%H:%M:%S')}")

                    # 统计在窗口内的数据点
                    in_window_count = sum(1 for ts in timestamps if ts >= cutoff_time)
                    print(f"  窗口内数据点: {in_window_count}")

                    # 显示最近几个数据点
                    print(f"  最近5个数据点:")
                    for i, (ts, vol) in enumerate(zip(list(timestamps)[-5:], list(volumes)[-5:])):
                        print(f"    {ts.strftime('%H:%M:%S')}: {vol}")
                else:
                    print("  无历史数据")
            else:
                print("  未初始化")

        except Exception as e:
            print(f"❌ 调试滑动窗口异常: {e}")


__all__ = ["solution"]